﻿using System;
using System.Collections.Generic;

namespace MetaSharp.Transformation
{
    // Set is an immutable hashtable, where you may have multiple keys with the same hashcode.
    // Keys should implement GetHashCode() and Equals(), or be unique by reference.
    class Set<TKey, TValue>
    {
        public static readonly Set<TKey, TValue> Empty = new Set<TKey, TValue>(null);

        private Bucket _root;

        public Bucket Root
        {
            get { return _root; }
        }

        private Set(Bucket root)
        {
            _root = root;
        }

        public Set<TKey, TValue> Add(TKey key, TValue value)
        {
            var hash = key.GetHashCode();
            var root = Set<TKey, TValue>.Add(
                _root,
                hash,
                key,
                value);

            return new Set<TKey, TValue>(root);
        }

        public IEnumerable<TValue> this[TKey key]
        {
            get { return Get(key); }
        }

        public IEnumerable<TValue> Get(TKey key)
        {
            var hash = key.GetHashCode();
            var found = Set<TKey, TValue>.Get(
                hash,
                key,
                _root);

            if (found == null)
                yield break;

            foreach (var node in found)
                yield return node.Value;
        }

        private static IEnumerable<Node> Get(int hash, TKey key, Bucket bucket)
        {
            if (bucket == null)
                return new Node[0];

            var comparison = bucket.GetHashCode().CompareTo(hash);
            if (comparison == 0)
            {
                var result = new List<Node>();
                var node = bucket.Root;
                while (node != null)
                {
                    if (object.Equals(node.Key, key))
                        result.Add(node);
                    node = node.Next;
                }

                return result;
            }

            if (comparison > 0)
                return Get(hash, key, bucket.Left);
            else
                return Get(hash, key, bucket.Right);
        }

        private static Bucket Add(Bucket bucket, int hash, TKey key, TValue value)
        {
            if (bucket == null)
                return new Bucket(hash, new Node(key, value, null), null, null);

            var comparison = bucket.GetHashCode().CompareTo(hash);
            if (comparison == 0)
            {
                return bucket.Add(key, value);
            }

            var l = bucket.Left;
            var r = bucket.Right;

            if (comparison > 0)
                l = Add(l, hash, key, value);
            else
                r = Add(r, hash, key, value);

            var n = bucket.SetBranches(l, r);
            var lh = n.Left == null ? 0 : n.Left.Height;
            var rh = n.Right == null ? 0 : n.Right.Height;
            var b = lh - rh;

            if (Math.Abs(b) == 2) // 2 or -2 means unbalanced
            {
                if (b == 2) // L
                {
                    var llh = n.Left.Left == null ? 0 : n.Left.Left.Height;
                    var lrh = n.Left.Right == null ? 0 : n.Left.Right.Height;
                    var lb = llh - lrh;
                    if (lb == 1) // LL
                    {
                        // rotate right
                        n = RotateRight(n);
                    }
                    else // LR
                    {
                        // rotate left
                        // rotate right
                        l = RotateLeft(l);
                        n = bucket.SetBranches(l, r);
                        //n = new Node<TKey, TValue>(n.Key, n.Value, l, r);
                        n = RotateRight(n);
                    }
                }
                else // R
                {
                    var rlh = n.Right.Left == null ? 0 : n.Right.Left.Height;
                    var rrh = n.Right.Right == null ? 0 : n.Right.Right.Height;
                    var rb = rlh - rrh;
                    if (rb == 1) // RL
                    {
                        // rotate right
                        // rotate left
                        r = RotateRight(r);
                        n = bucket.SetBranches(l, r);
                        //n = new Node<TKey, TValue>(n.Key, n.Value, l, r);
                        n = RotateLeft(n);
                    }
                    else // RR
                    {
                        // rotate left
                        n = RotateLeft(n);
                    }
                }
            }

            return n;
        }

        private static Bucket RotateRight(Bucket bucket)
        {
            //       (5)            4     
            //       / \           / \    
            //      4   D         /   \   
            //     / \           3     5  
            //    3   C    -->  / \   / \ 
            //   / \           A   B C   D
            //  A   B                     

            var L = bucket.Left.Left;
            //var R = new Node<TKey, TValue>(node.Key, node.Value, node.Left.Right, node.Right);
            //var N = new Node<TKey, TValue>(node.Left.Key, node.Left.Value, L, R);
            var R = bucket.SetBranches(bucket.Left.Right, bucket.Right);
            var B = bucket.Left.SetBranches(L, R);
            return B;
        }

        private static Bucket RotateLeft(Bucket bucket)
        {
            //    (3)               4     
            //    / \              / \    
            //   A   4            /   \   
            //      / \          3     5  
            //     B   5   -->  / \   / \ 
            //        / \      A   B C   D
            //       C   D                

            //var L = new Node<TKey, TValue>(node.Key, node.Value, node.Left, node.Right.Left);
            var L = bucket.SetBranches(bucket.Left, bucket.Right.Left);
            var R = bucket.Right.Right;
            //var N = new Node<TKey, TValue>(node.Right.Key, node.Right.Value, L, R);
            var B = bucket.Right.SetBranches(L, R);
            return B;
        }

        public class Bucket
        {
            private int _hash;
            private int _height;
            private Node _root;
            private Bucket _left;
            private Bucket _right;

            public Node Root { get { return _root; } }
            public Bucket Left { get { return _left; } }
            public Bucket Right { get { return _right; } }
            public int Height { get { return _height; } }

            public Bucket(int hash, Node root, Bucket left, Bucket right)
            {
                _hash = hash;
                _right = right;
                _left = left;
                _root = root;

                var l = left == null ? 0 : left.Height;
                var r = right == null ? 0 : right.Height;
                this._height = Math.Max(l, r) + 1;
            }

            public Bucket Add(TKey key, TValue value)
            {
                var root = Remove(key, _root);
                return new Bucket(
                    _hash,
                    new Node(key, value, root),
                    _left,
                    _right);
            }

            private Node Remove(TKey key, Node node)
            {
                var same = true;
                return Remove(key, _root, out same);
            }

            private Node Remove(TKey key, Node node, out bool same)
            {
                if (node == null)
                {
                    same = true;
                    return null;
                }

                var next = Remove(key, node.Next, out same);
                if(!same)
                {
                    same = false;
                    return new Node(node.Key, node.Value, next);
                }

                if (object.Equals(node.Key, key))
                {
                    same = false;
                    return next;
                }

                same = true;
                return node;
            }

            public Bucket SetBranches(Bucket left, Bucket right)
            {
                return new Bucket(
                    _hash,
                    _root,
                    left,
                    right);
            }

            public override int GetHashCode()
            {
                return _hash;
            }
        }

        public class Node
        {
            private TKey _key;
            private TValue _value;
            private Node _next;

            public TKey Key { get { return _key; } }
            public TValue Value { get { return _value; } }
            public Node Next { get { return _next; } }

            public Node(TKey key, TValue value, Node next)
            {
                _key = key;
                _value = value;
                _next = next;
            }
        }
    }
}
